"""
B卷
1.	把inceptionnet-v3预训练模型，作为主干网络，将5种农作物图像数据，映射为2048维向量。
按照下述要求，完成相应操作（30分）
"""
import os
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import cv2 as cv
import seaborn as sns

# ①	定义数据参数，包括：图像文件夹、图像特征向量保存地址、inception-v3模型路径和模型参数等
INPUT_PLACEHOLDER_NAME = 'DecodeJpeg/contents:0'
OUTPUT_TENSOR_NAME = 'pool_3/_reshape:0'
IMG_DIR_TRAIN = 'data/agricultrue/train'
IMG_DIR_TEST = 'data/agricultrue/test'
MODEL_PATH = 'data/inceptionnet_v3/tensorflow_inception_graph.pb'
VER = 'v1.0'
FEATURE_SAVE_DIR = os.path.join('vec_output', VER)
os.makedirs(FEATURE_SAVE_DIR, exist_ok=True)
FILE_NAME = os.path.basename(__file__)
LOG_DIR = os.path.join('_log', FILE_NAME, VER)


def sep(label=''):
    print('-' * 32, label, '-' * 32)


# ②	创建计算图，加载inception-v3预训练模型，返回数据输入张量和瓶颈层输出张量
sep('Import model')
graphDef = tf.GraphDef()
with open(MODEL_PATH, 'rb') as f:
    model_bin = f.read()
graphDef.ParseFromString(model_bin)
with tf.Session() as sess:
    input_placeholder, output_tensor = tf.import_graph_def(graphDef, return_elements=[
        INPUT_PLACEHOLDER_NAME,
        OUTPUT_TENSOR_NAME
    ])
    # with tf.summary.FileWriter(LOG_DIR) as fw:
    #     fw.add_graph(sess.graph)

    # ③	开启会话，读取所有的图像，将图像映射的特征向量保存在相应地址
    sep('Process images to extract feature vector')
    sub_dirs = os.listdir(IMG_DIR_TRAIN)
    cnt = 0
    skip = 0
    for sub_dir in sub_dirs:
        sub_dir_path = os.path.join(IMG_DIR_TRAIN, sub_dir)
        output_sub_dir = os.path.join(FEATURE_SAVE_DIR, sub_dir)
        os.makedirs(output_sub_dir, exist_ok=True)
        files = os.listdir(sub_dir_path)
        for file_name in files:
            if (cnt + skip) % 25 == 0:
                print(f'Handled {(cnt + skip)} pictures.')
            output_file_path = os.path.join(output_sub_dir, file_name + '.txt')
            if os.path.exists(output_file_path):
                skip += 1
                continue

            file_path = os.path.join(sub_dir_path, file_name)
            with open(file_path, 'rb') as f:
                img_bin = f.read()
            feature_vec = sess.run(output_tensor, feed_dict={input_placeholder: img_bin})
            feature_vec = np.squeeze(feature_vec, axis=0)
            np.savetxt(output_file_path, feature_vec)
            cnt += 1
    print(f'Picture processing over. {cnt} pictures are processed and {skip} pictures used cache.')


    # ④	给定的3张测试图像，计算它们之间的相似度（3张测试图像，如下图所示）
    sep('Calculate cosine correlation of pictures')
    def cos(a, b):
        ab = np.sum(a * b)
        ra = np.linalg.norm(a)
        rb = np.linalg.norm(b)
        cos = ab / (ra * rb + 1e-20)
        return cos


    PIC_SIDE = 299
    pics_to_test = [
        'corn_008.jpg',
        'millet_007.jpg',
        'millet_014.jpg'
    ]
    TEST_LEN = len(pics_to_test)
    pics_to_test = sorted(pics_to_test)  # Sort by name !
    pics_path = [os.path.join(IMG_DIR_TEST, name) for name in pics_to_test]
    vecs_to_test = []
    for path in pics_path:
        img = cv.imread(path, cv.IMREAD_COLOR)
        img = cv.resize(img, (PIC_SIDE, PIC_SIDE))
        img = img.reshape(1, -1)
        img = np.squeeze(img, axis=0).astype(np.float32)  # ATTENTION Must float32 rather than uint8
        vecs_to_test.append(img)
    cos_mat = np.zeros([TEST_LEN, TEST_LEN])
    for i in range(TEST_LEN):
        for j in range(TEST_LEN):
            cos_ij = cos(vecs_to_test[i], vecs_to_test[j])
            cos_mat[i, j] = cos_ij
    print(cos_mat)

    # ⑤	计算这3张图像映射的特征向量之间的相似度
    sep('Calculate cosine correlation of feature vectors')
    feature_vecs_to_test = []
    for path in pics_path:
        with open(path, 'rb') as f:
            img_bin = f.read()
        feature_vec = sess.run(output_tensor, feed_dict={input_placeholder: img_bin})
        feature_vec = np.squeeze(feature_vec, axis=0)
        feature_vecs_to_test.append(feature_vec)
    feature_cos_mat = np.zeros([TEST_LEN, TEST_LEN])
    for i in range(TEST_LEN):
        for j in range(TEST_LEN):
            cos_ij = cos(feature_vecs_to_test[i], feature_vecs_to_test[j])
            feature_cos_mat[i, j] = cos_ij
    print(feature_cos_mat)

    # ⑥	通过图示，比较图像映射（即特征提取）前后相似度变化情况
    sep('Show them')
    spr = 2
    spc = 3
    spn = 0
    plt.figure(figsize=[10, 5])
    for i, name in enumerate(pics_to_test):
        spn += 1
        plt.subplot(spr, spc, spn)
        plt.title(name)
        plt.axis('off')
        img = plt.imread(pics_path[i])
        plt.imshow(img)

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('cos between pictures')
    sns.heatmap(cos_mat, annot=True)

    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title('cos between features')
    sns.heatmap(feature_cos_mat, annot=True)

    print('Please check the plotting window.')
    plt.show()
    print('Over')
